"""
ASSUMES UTF-8
"""

import codecs
import re
from utils import util

from nltk_lite import tokenize
from lib.probability import FreqDist

import sys
sys.stdout = codecs.lookup('cp737')[-1](sys.stdout)



def countWordsTXT(filename_in, words = FreqDist(), check=lambda x1:x1.isalpha(), transform=None):
    file = codecs.open(filename_in,'r','utf-8')
    contents = file.read()
    pattern = re.compile(r'\w+', re.UNICODE) 
    for word in tokenize.regexp(contents, pattern):
        if re.match(pattern,word) and check(word):
            if transform:
                word = transform(word)
            words.inc(word)
    return words

def countWordsTXTold(filename_in, words = FreqDist()):
    file = codecs.open(filename_in,'r','utf-8')
    contents = file.read()
    pattern = re.compile(r'\w+', re.UNICODE) 
    for word in tokenize.regexp(contents, pattern):
        if re.match(pattern,word) and word.isalpha():
            words.inc(word)
    return words

def countWordsXML(filename, words=FreqDist()):
    #file = codecs.open(filename,'r','utf-8')
    import xml.parsers.expat
    file = open(filename, 'r')
    p = xml.parsers.expat.ParserCreate('utf-8')
    pattern = re.compile(r'\w+', re.UNICODE) 
    def char_data(data):
        #print encode(data,'ascii')
        #print repr(data)
        if re.match(pattern,data) and data.isalpha():
            #print repr(data)
            words.inc(data)
    p.CharacterDataHandler = char_data
    p.returns_unicode=1
    p.ParseFile(file)
    
    return words

def getHandler(filename):
    ext = os.path.splitext(filename)[1].upper()[1:]
    subclass = "countWords%s" % ext
    module = sys.modules[M.__module__]
    func = hasattr(module, subclass) and getattr(module, subclass) or countWordsTXT
    return func
    

import os
def countWordsinDirectory(directory):
    words = FreqDist()
    for filename in util.all_files(directory, '*'):
        print 'Parsing file %s' % filename
        func = getHandler(filename)
        words = func(filename, words)
                    
    return words

class M:
    pass


def countWordsXMLinDirectory(directory):
    words = FreqDist()
    for file in util.all_files(directory, '*.xml'):
        print 'Parsing file %s' % file
        words = countWordsXML(file, words)
    return words


def read_words(filename, words=FreqDist()):
    """reads the words in filename 
    in the form of:
    word - count - (frequency)
    and return a FreqDist object"""
    file = codecs.open(filename, mode = 'r', encoding='utf-8')
    for line in file:
        [word, count, freq]=line.split('-')
        words.inc(word, count)
    return words        

        
            
        
if __name__ == '__main__':
    trans = lambda x1:x1.lower()
    words = countWordsTXT('hellas_old.txt', check=lambda x1:True, transform=trans)
    words = countWordsTXT('hellas05-07.txt', words, check=lambda x1:True, transform=trans)    
    words = countWordsTXT('hellas06-07.txt', words, check=lambda x1:True, transform=trans)    
    print words.N()
    file = codecs.open('words-grkl.txt', mode = 'wb', encoding='utf-8') 
    for word in words.sorted_samples():
        file.write(word)
        file.write(' - ')
        file.write(str(words.count(word)))
        file.write('\r\n')        
    

if __name__=='__main__1':
    print sys.argv
    root, out = None, None
    try:
        root = sys.argv[1]
        out = sys.argv[2]
    except:
        pass

    if(root):
        words = FreqDist()
        if os.path.isdir(root):
            words = countWordsinDirectory(root)
        else:
            func = getHandler(root)
            words = func(root)
        if(out):
            util.print_words(words, out)
        else:
            util.print_words(words)               
                       

        
    

    
        
        
                      
    

    

    
